package io.kiki.sba.registry.common.model.slot;

import com.fasterxml.jackson.annotation.JsonIgnore;
import io.kiki.sba.registry.exception.SofaRegistrySlotsException;
import io.kiki.sba.registry.util.JsonUtils;
import lombok.Getter;
import org.springframework.util.CollectionUtils;

import java.io.Serializable;
import java.util.*;


public final class Slots implements Serializable {
    public static final Slots INIT = new Slots(-1, Collections.emptyList());
    @Getter
    private final long epoch;
    private final Map<Integer, Slot> slotIdToSlotMap;

    private Slots(long epoch, Map<Integer, Slot> slotIdToSlotMap) {
        this.epoch = epoch;
        if (!CollectionUtils.isEmpty(slotIdToSlotMap)) {
            this.slotIdToSlotMap = new TreeMap<>(slotIdToSlotMap);
        } else {
            this.slotIdToSlotMap = new TreeMap<>();
        }
    }

    public Slots(long epoch, final Collection<Slot> slotIdToSlotMap) {
        this.epoch = epoch;
        SortedMap<Integer, Slot> slotMap = new TreeMap<>();
        if (!CollectionUtils.isEmpty(slotIdToSlotMap)) {
            slotIdToSlotMap.forEach(slot -> {
                Slot exist = slotMap.putIfAbsent(slot.getId(), slot);
                if (exist != null) {
                    throw new SofaRegistrySlotsException("dup slot when construct slot table: " + JsonUtils.writeValueAsString(slotIdToSlotMap));
                }
            });
        }
        this.slotIdToSlotMap = slotMap;
    }

    public List<Slot> getSlotIdToSlotMap() {
        return new ArrayList<>(slotIdToSlotMap.values());
    }

    @JsonIgnore
    public Map<Integer, Slot> getSlotMap() {
        return new HashMap<>(slotIdToSlotMap);
    }

    @JsonIgnore
    public Set<Integer> getSlotIds() {
        return new TreeSet<>(slotIdToSlotMap.keySet());
    }

    public Slot getSlot(int slotId) {
        return slotIdToSlotMap.get(slotId);
    }

    public Map<Integer, String> slotLeaders() {
        Map<Integer, String> ret = new TreeMap<>();
        for (Map.Entry<Integer, Slot> e : slotIdToSlotMap.entrySet()) {
            ret.put(e.getKey(), e.getValue().getLeaderDataStoreNodeId());
        }
        return ret;
    }


    public List<DataStoreNode> transfer(String targetDataNode, boolean ignoreFollowers) {
        Map<String, List<Integer>> leadersMap = new HashMap<>();
        Map<String, List<Integer>> followersMap = new HashMap<>();
        for (Slot slot : slotIdToSlotMap.values()) {
            if (targetDataNode == null) {
                List<Integer> leaders = leadersMap.computeIfAbsent(slot.getLeaderDataStoreNodeId(), k -> new ArrayList<>());
                leaders.add(slot.getId());
                if (!ignoreFollowers) {
                    for (String follower : slot.getFollowerDataStoreNodeIdSet()) {
                        List<Integer> followers = followersMap.computeIfAbsent(follower, k -> new ArrayList<>());
                        followers.add(slot.getId());
                    }
                }
                continue;
            }
            if (!ignoreFollowers && slot.getFollowerDataStoreNodeIdSet().contains(targetDataNode)) {
                List<Integer> followers = followersMap.computeIfAbsent(targetDataNode, k -> new ArrayList<>());
                followers.add(slot.getId());
            }
            if (targetDataNode.equals(slot.getLeaderDataStoreNodeId())) {
                List<Integer> leaders = leadersMap.computeIfAbsent(slot.getLeaderDataStoreNodeId(), k -> new ArrayList<>());
                leaders.add(slot.getId());
            }
        }
        Map<String, DataStoreNode> dataNodeSlotMap = new HashMap<>(leadersMap.size());
        for (Map.Entry<String, List<Integer>> entry : leadersMap.entrySet()) {
            DataStoreNode dataStoreNode = dataNodeSlotMap.computeIfAbsent(entry.getKey(), k -> new DataStoreNode(k));
            Collections.sort(entry.getValue());
            dataStoreNode.addLeaderSlotId(entry.getValue());
        }
        for (Map.Entry<String, List<Integer>> entry : followersMap.entrySet()) {
            DataStoreNode dataStoreNode = dataNodeSlotMap.computeIfAbsent(entry.getKey(), k -> new DataStoreNode(k));
            Collections.sort(entry.getValue());
            dataStoreNode.addFollowerSlotId(entry.getValue());
        }
        return new ArrayList<>(dataNodeSlotMap.values());
    }

    public Set<String> getDataServers() {
        final Set<String> set = new TreeSet<>();
        slotIdToSlotMap.values().forEach(slot -> {
            set.add(slot.getLeaderDataStoreNodeId());
            set.addAll(slot.getFollowerDataStoreNodeIdSet());
        });
        return set;
    }

    public Slots filter(Set<Integer> slotIds) {
        if (slotIdToSlotMap.isEmpty()) {
            return this;
        }
        final Map<Integer, Slot> map = new TreeMap<>();
        for (Integer slotId : slotIds) {
            map.put(slotId, slotIdToSlotMap.get(slotId));
        }
        return new Slots(epoch, map);
    }

    public Slots filter(String ip) {
        if (slotIdToSlotMap.isEmpty()) {
            return this;
        }
        final Map<Integer, Slot> map = new TreeMap<>();
        slotIdToSlotMap.forEach((k, v) -> {
            if (v.getLeaderDataStoreNodeId().equals(ip) || v.getFollowerDataStoreNodeIdSet().contains(ip)) {
                map.put(k, v);
            }
        });
        return new Slots(epoch, map);
    }

    public Slots filterLeaderInfo() {
        if (slotIdToSlotMap.isEmpty()) {
            return this;
        }
        final Map<Integer, Slot> map = new TreeMap<>();
        slotIdToSlotMap.forEach((k, v) -> {
            // filter followers
            map.put(k, new Slot(v.getId(), v.getLeaderDataStoreNodeId(), v.getLeaderEpoch(), Collections.emptyList()));
        });
        return new Slots(epoch, map);
    }

    public int getLeaderNum(String dataServerIp) {
        return (int) slotIdToSlotMap.values().stream().filter(s -> dataServerIp.equals(s.getLeaderDataStoreNodeId())).count();
    }

    public int getSlotNum() {
        return slotIdToSlotMap.size();
    }

    public int getFollowerNum(String dataServerIp) {
        return (int) slotIdToSlotMap.values().stream().filter(s -> s.getFollowerDataStoreNodeIdSet().contains(dataServerIp)).count();
    }

    public void assertSlotLessThan(Slots slots) {
        this.slotIdToSlotMap.values().forEach(slot -> {
            Slot newSlot = slots.getSlot(slot.getId());
            if (newSlot == null) {
                return;
            }
            if (slot.getLeaderEpoch() > newSlot.getLeaderEpoch()) {
                throw new RuntimeException(String.format("not expect Slot.LeaderEpoch, %s, %s", slot, newSlot));
            }
        });
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        Slots slots = (Slots) o;
        return epoch == slots.epoch && this.slotIdToSlotMap.equals(slots.slotIdToSlotMap);
    }

    @Override
    public int hashCode() {
        return Objects.hash(epoch, slotIdToSlotMap);
    }

}
